import torch
import torch.nn as nn
import torch.nn.functional as F

class UNetDown(nn.Module):
    """U-Net下采样模块"""
    def __init__(self, in_channels, out_channels, normalize=True, dropout=0.0):
        super(UNetDown, self).__init__()
        
        # 使用卷积层进行下采样，步长为2
        layers = [nn.Conv2d(in_channels, out_channels, 4, stride=2, padding=1, bias=False)]
        
        # 是否使用批归一化
        if normalize:
            layers.append(nn.BatchNorm2d(out_channels))
        
        # 使用LeakyReLU激活函数
        layers.append(nn.LeakyReLU(0.2))
        
        # 是否使用Dropout
        if dropout:
            layers.append(nn.Dropout(dropout))
        
        self.model = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.model(x)

class UNetUp(nn.Module):
    """U-Net上采样模块"""
    def __init__(self, in_channels, out_channels, dropout=0.0):
        super(UNetUp, self).__init__()
        
        # 使用转置卷积进行上采样
        layers = [
            nn.ConvTranspose2d(in_channels, out_channels, 4, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        ]
        
        # 是否使用Dropout
        if dropout:
            layers.append(nn.Dropout(dropout))
        
        self.model = nn.Sequential(*layers)
    
    def forward(self, x, skip_input):
        # 将上采样结果与跳跃连接的特征图拼接
        x = self.model(x)
        x = torch.cat((x, skip_input), 1)
        return x

class ColorizationGenerator(nn.Module):
    """图像着色生成器模型"""
    def __init__(self, in_channels=1, out_channels=2):
        super(ColorizationGenerator, self).__init__()
        
        # 输入层
        self.down1 = UNetDown(in_channels, 64, normalize=False)
        # 下采样层
        self.down2 = UNetDown(64, 128)
        self.down3 = UNetDown(128, 256)
        self.down4 = UNetDown(256, 512, dropout=0.5)
        self.down5 = UNetDown(512, 512, dropout=0.5)
        self.down6 = UNetDown(512, 512, dropout=0.5)
        self.down7 = UNetDown(512, 512, dropout=0.5)
        self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5)
        
        # 上采样层
        self.up1 = UNetUp(512, 512, dropout=0.5)
        self.up2 = UNetUp(1024, 512, dropout=0.5)
        self.up3 = UNetUp(1024, 512, dropout=0.5)
        self.up4 = UNetUp(1024, 512, dropout=0.5)
        self.up5 = UNetUp(1024, 256)
        self.up6 = UNetUp(512, 128)
        self.up7 = UNetUp(256, 64)
        
        # 输出层
        self.final = nn.Sequential(
            nn.ConvTranspose2d(128, out_channels, 4, stride=2, padding=1),
            nn.Tanh()
        )
    
    def forward(self, x):
        # 下采样路径
        d1 = self.down1(x)
        d2 = self.down2(d1)
        d3 = self.down3(d2)
        d4 = self.down4(d3)
        d5 = self.down5(d4)
        d6 = self.down6(d5)
        d7 = self.down7(d6)
        d8 = self.down8(d7)
        
        # 上采样路径
        u1 = self.up1(d8, d7)
        u2 = self.up2(u1, d6)
        u3 = self.up3(u2, d5)
        u4 = self.up4(u3, d4)
        u5 = self.up5(u4, d3)
        u6 = self.up6(u5, d2)
        u7 = self.up7(u6, d1)
        
        # 输出
        return self.final(u7)

class ColorizationDiscriminator(nn.Module):
    """图像着色判别器模型"""
    def __init__(self, in_channels=3):
        super(ColorizationDiscriminator, self).__init__()
        
        # 判别器网络结构
        self.model = nn.Sequential(
            # 第一层不使用BatchNorm
            nn.Conv2d(in_channels, 64, 4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            
            # 后续层
            nn.Conv2d(64, 128, 4, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(128, 256, 4, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(256, 512, 4, stride=1, padding=1),
            nn.BatchNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            # 输出层
            nn.Conv2d(512, 1, 4, stride=1, padding=1)
        )
    
    def forward(self, img_a, img_b):
        # 将灰度图和彩色图拼接
        img_input = torch.cat((img_a, img_b), 1)
        return self.model(img_input) 